__author__ = 'carlxie'

import numpy as np

from layer.Layer import Layer
from layer.LayerType import LayerType
from layer.ConvLayer import ConvLayer


class MaxPoolLayer(Layer):

    def __init__(self,preLayer,F=2,S=1):
        assert (preLayer.width - F) % S == 0
        assert (preLayer.height - F) % S == 0
        super(MaxPoolLayer,self).__init__(
            preLayer,
            LayerType.MAX_POOL_TYPE,
            (preLayer.width - F) / S + 1,
            (preLayer.height - F) / S + 1,
            preLayer.depth
        )
        self.kernelSize = F
        self.stride = S

    def forward_prop(self):
        self.poolIdx2MaxIdx = {}
        input_data = self.preLayer.output[:,:,:]
        shape = self.output.shape
        for depthIdx in range(shape[2]):
            for heightIdx in range(shape[1]):
                for widthIdx in range(shape[0]):
                    self.output[widthIdx,heightIdx,depthIdx] = self.max_pool(input_data,depthIdx,heightIdx,widthIdx)
        self.score = self.output

    def max_pool(self,data,dIdx,hIdx,wIdx):
        w = wIdx * self.stride
        h = hIdx * self.stride
        maxVal = -999999999
        for i in range(self.kernelSize):
            for j in range(self.kernelSize):
               if maxVal < data[w+i,h+j,dIdx]:
                   maxVal = data[w+i,h+j,dIdx]
                   self.poolIdx2MaxIdx[(wIdx,hIdx,dIdx)] = (w+i,h+j,dIdx)
        return maxVal

    def back_prop(self):
        self.preLayer.delta = np.zeros(self.preLayer.output.shape)
        for key in self.poolIdx2MaxIdx.keys():
            d = self.delta[key]
            k = self.poolIdx2MaxIdx[key]
            self.preLayer.delta[k] = d


if __name__ == "__main__":
    b = np.array([
        [2, 1, 1, 1, 1],
        [2, 0, 2, 0, 1],
        [1, 0, 0, 1, 2],
        [2, 0, 2, 1, 0],
        [0, 2, 2, 2, 1]])
    c = np.array([
        [1, 1, 0, 2, 1],
        [2, 0, 1, 1, 0],
        [2, 2, 1, 1, 2],
        [0, 1, 1, 2, 1],
        [2, 1, 2, 2, 2]])
    a = np.array([
        [0, 0, 2, 2, 0],
        [0, 1, 1, 1, 1],
        [1, 2, 1, 2, 0],
        [2, 0, 2, 2, 0],
        [0, 0, 2, 1, 1]
    ])
    k1 = np.array([
        [ 0,  1, -1],
        [-1,  0,  0],
        [ 0,  0, -1]
    ])
    k2 = np.array([
        [-1,  1,  1],
        [-1, -1, -1],
        [ 1,  0,  1]
    ])
    k3 = np.array([
        [ 0,  1,  1],
        [ 1, -1,  0],
        [ 1,  0, -1]
    ])
    weights = np.zeros((3,3,3))
    weights[:,:,0] = k1
    weights[:,:,1] = k2
    weights[:,:,2] = k3

    input_data = np.zeros((5,5,3))
    input_data[:,:,0] = a
    input_data[:,:,1] = b
    input_data[:,:,2] = c


    biases = np.empty((1,1,1))
    biases[:,:,0] = 1

    layer = Layer(None,LayerType.INPUT_TYPE,5,5,3)
    layer.output = input_data

    convLayer = ConvLayer(layer,1,3,2,1)
    convLayer.weights = weights
    convLayer.biases = biases
    convLayer.forward_prop()

    print convLayer.output[:,:,0]

    maxLayer = MaxPoolLayer(convLayer)
    maxLayer.forward_prop()

    print maxLayer.output[:,:,0]